#include "dev/disk.h"
#include "tools/log.h"
#include "tools/klib.h"
#include "comm/cpu_instr.h"
#include "comm/boot_info.h"
#include "dev/dev.h"
#include "cpu/irq.h"

static disk_t disk_table[DISK_TABLE_SIZE];
static mutex_t mutex;
static semaphore_t sema;
// 用于控制磁盘中断程序相应，
// 因为boot loader中的磁盘加载程序，所以导致中断应该堵车了
// 我们只需要磁盘响应disk_read和disk_write的读写请求
static int disk_on_op = 0;

static void disk_send_cmd(disk_t * disk, 
                        uint16_t start_sector, 
                        uint32_t sector_count, 
                        int cmd)
{
    outb(DISK_DRIVE(disk), DISK_DRIVE_BASE | disk->drive);
    // 高八位
    outb(DISK_SECTOR_COUNT(disk), (uint8_t)(sector_count >> 8));
    outb(DISK_LBA_L0(disk), (uint8_t)(start_sector >> 24));   
    outb(DISK_LBA_MID(disk), 0);
    outb(DISK_LBA_HI(disk),0); 

    // 低八位
    outb(DISK_SECTOR_COUNT(disk),(uint8_t)(sector_count));
    outb(DISK_LBA_L0(disk),(uint8_t)(start_sector));
    outb(DISK_LBA_MID(disk),(uint8_t)(start_sector>>8));
    outb(DISK_LBA_HI(disk),(uint8_t)(start_sector>>16));

    outb(DISK_CMD(disk), (uint8_t)cmd);
}

static void disk_read_data(disk_t * disk, void * buff, int size)
{
    uint16_t *c = (uint16_t *)buff;
    for(int i=0;i<size/2;i++)
    {
        *c++ = inw(DISK_DATA(disk));
    }
}

static void disk_write_data(disk_t * disk, void * buff, int size)
{
    uint16_t *c = (uint16_t *)buff;
    for(int i=0;i<size/2;i++)
    {
        outw(DISK_DATA(disk), *c++);
    }
}

static int disk_wait_data(disk_t * disk)
{
    uint8_t status;
    do{
        status = inb(DISK_STATUS(disk));
        if((status & (DISK_STATUS_ERR|DISK_STATUS_DRQ|DISK_STATUS_DF|DISK_STATUS_BSY)) != DISK_STATUS_BSY)
        {
            break;
        }
    }while(1);  

    return (status & DISK_STATUS_ERR) ? -1 : 0;
}

static void print_disk_info(disk_t * disk)
{
    log_printf("Disk name: %s", disk->name);
    log_printf("    base port: %x", disk->port_base);
    log_printf("    sector count: %d, total size: %d m", 
                disk->sector_count, 
                disk->sector_count * disk->sector_size / 1024 / 1024);

    for(int i=0; i<DISK_PRIMARY_PART_COUNT; i++)
    {
        partinfo_t *part = disk->partinfo + i;
        if(part->type != FS_INVALID)
        {
            log_printf("    part: %d",i);
            log_printf("            type: %d",part->type);
            log_printf("            name: %s",part->name);
            log_printf("            start sector: %d",part->start_sector);
            log_printf("            total sector: %d",part->total_sector_count);
        }
    }
}

static int detect_partinfo(disk_t * disk)
{
    mbr_t mbr;

    disk_send_cmd(disk, 0, 1, DISK_CMD_READ);
    int err = disk_wait_data(disk);
    if(err < 0)
    {
        log_printf("read mbr failed.");
        return err;
    }

    // 磁盘上的信息读入内存 
    disk_read_data(disk, &mbr, sizeof(mbr_t));
    part_item_t *item = mbr.part_item;
    partinfo_t *part_info = disk->partinfo + 1;
    for(int i=0; i<MBR_PRIMARY_PART_COUNT; i++, item++, part_info++)
    {
        part_info->type = item->system_id;
        if(part_info->type == FS_INVALID)
        {
            part_info->start_sector = part_info->total_sector_count = 0;
            part_info->disk = 0;
            kernel_memset(part_info->name, 0, DISK_PART_NAME_SIZE);
        }
        else
        {
            kernel_sprintf(part_info->name, "%s%d", disk->name, i+1);
            part_info->disk = disk;
            part_info->start_sector = item->relative_sector;
            part_info->total_sector_count = item->total_sectors;
        }
    }
}

static int identify_disk(disk_t * disk)
{
    // 初始化磁盘
    disk_send_cmd(disk, 0, 0, DISK_CMD_IDENTIFY);
    int err = inb(DISK_STATUS(disk));
    if(err == 0)
    {
        log_printf("disk does not exist", disk->name);
        return -1;
    }

    err = disk_wait_data(disk);
    if(err < 0)
    {
        log_printf("disk[%s] read failed", disk->name);
        return err;
    }

    uint16_t buff[256];
    disk_read_data(disk, buff, sizeof(buff));
    disk->sector_count = *(uint32_t *)(buff + 100);
    disk->sector_size = SECTOR_SIZE;

    // 取第一个多余的分区，用于管理磁盘
    partinfo_t *part = disk->partinfo + 0;
    part->disk = disk;
    part->start_sector = 0;
    part->total_sector_count = disk->sector_count;
    part->type = FS_INVALID;        // 由于这个只用于管理磁盘，不是文件系统分区，所以不可用
    kernel_sprintf(part->name, "%s%d", disk->name, 0); 

    // 分区信息检查
    detect_partinfo(disk);

    return 0;
}

void disk_init(void)
{
    log_printf("Check disk......");
    kernel_memset((void *)disk_table, 0, sizeof(disk_table));
    mutex_init(&mutex);
    sema_init(&sema, 0);
    for(int i=0; i<DISK_PER_CHANNEL; i++)
    {
        disk_t * disk = disk_table + i;
        kernel_sprintf(disk->name, "sd%c", i + 'a');
        disk->drive = (i==0) ? DISK_MASTER : DISK_SLAVE;
        disk->port_base = DISK_IOBASE_PRIMARY;
        disk->mutex = &mutex;
        disk->operate_sema = &sema;
        // 磁盘检测
        int err = identify_disk(disk);
        if(err == 0) print_disk_info(disk);
    }
}

void do_handler_disk_primary(exception_frame_t * frame)
{
    irq_send_eoi(IRQ0X2E_HARDDISK_PRIMARY);
    // 准备数据，准备完成发送信号
    if(disk_on_op && curr_task_get())
    {
        sema_notify(&sema);
    }
}

int disk_open(device_t * dev)
{
    // 0xa0-- a:磁盘编号，0:分区
    int disk_idx = (dev->minor >> 4) - 0xa;
    int part_idx = dev->minor & 0xF;

    if(disk_idx >= DISK_TABLE_SIZE || part_idx >= DISK_PRIMARY_PART_COUNT)
    {
        log_printf("disk minor error. device: sd%x", dev->minor);
        return -1;
    }

    disk_t * disk = disk_table + disk_idx;
    if(!disk)
    {
        log_printf("disk not exist. device: sd%x", dev->minor);
        return -1;
    }

    partinfo_t * part = disk->partinfo + part_idx;
    if(part->total_sector_count == 0)
    {
        log_printf("part not exist. device: sd%x", dev->minor);
        return -1;
    }

    dev->data = part;

    irq_install(IRQ0X2E_HARDDISK_PRIMARY, (uint32_t)exception_handler_disk_primary);
    irq_enable(IRQ0X2E_HARDDISK_PRIMARY);
    return 0;
}

// 设备/LBA模式下相对起始扇区的扇区号/需要写入的数据/扇区数量
int disk_read(device_t * dev, int relative_sector, char * buff, int sector_count)
{   // part = dev.data;
    partinfo_t * part = (partinfo_t *)dev->data;
    if(!part)
    {
        log_printf("Get part info error. device: %d", dev->minor);
        return -1;
    }

    disk_t * disk = part->disk;
    if(!disk)
    {
        log_printf("No disk in PartInfo. device: %d", dev->minor);
    }

    mutex_lock(disk->mutex);
    // 发送读取命令，此时OS开始准备数据，
    disk_on_op = 1;
    disk_send_cmd(disk, part->start_sector + relative_sector, sector_count, DISK_CMD_READ);
    
    int read_sector_count = 0;
    for(; read_sector_count<sector_count; read_sector_count++, buff += disk->sector_size)
    {
        // 等待数据准备，由中断处理程序激活信号量
        if(curr_task_get()) sema_wait(disk->operate_sema);
        int err = disk_wait_data(disk);
        if(err < 0)
        {
            log_printf("read disk(%s) failed. start sector: %d, count: %d",
                disk->name, relative_sector, sector_count
            );
            break;
        }
        disk_read_data(disk, buff, disk->sector_size);
    }
    mutex_unlock(disk->mutex);
    return read_sector_count;
}

int disk_write(device_t * dev, int relative_sector, char * buff, int sector_count)
{
    partinfo_t * part = (partinfo_t *)dev->data;
    if(!part)
    {
        log_printf("Get part info error. device: %d", dev->minor);
        return -1;
    }

    disk_t * disk = part->disk;
    if(!disk)
    {
        log_printf("No disk in PartInfo. device: %d", dev->minor);
    }

    mutex_lock(disk->mutex);
    // 发送读取命令，此时OS开始准备数据，
    disk_on_op = 1;
    disk_send_cmd(disk, part->start_sector + relative_sector, sector_count, DISK_CMD_WRITE);
    
    int read_sector_count = 0;
    for(; read_sector_count<sector_count; read_sector_count++, buff += disk->sector_size)
    {
        disk_write_data(disk, buff, disk->sector_size);
        // 等待数据写入完成，由中断处理程序激活信号量
        if(curr_task_get()) sema_wait(disk->operate_sema);
        int err = disk_wait_data(disk);
        if(err < 0)
        {
            log_printf("write disk(%s) failed. start sector: %d, count: %d",
                disk->name, relative_sector, sector_count
            );
            break;
        }
    }
    mutex_unlock(disk->mutex);
    return read_sector_count;
}

int disk_control(device_t * dev, int cmd, int arg0, int arg1)
{
    return -1;
}

void disk_close(device_t * dev)
{

}

device_desc_t dev_disk_desc = 
{
    .name = "disk",
    .major = DEV_DISK,
    .open = disk_open,
    .read = disk_read,
    .write = disk_write,
    .control = disk_control,
    .close = disk_close,
};